#-*-coding:utf-8-*-#

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import container

os.chdir('PATH')

"""

Simulation results: Data_Fig4.zip

"""



dataframes = {}

nlist = [5,10,15,20,25,30,35,40,42,45,50]

for pnum in [2,3,4,5]:
    for i in nlist:
        dataframes[f'dr{i}_{pnum}'] = pd.read_csv('drought simulation AIDD t '+ str(i) +' number '+str(pnum)+'.csv',index_col=0)

nlist = [5,8,10,15,20,25,30,35,40,45,50]

for pnum in [2,3,4,5]:
    for i in nlist:
        dataframes[f'wf{i}_{pnum}'] = pd.read_csv('wildfire simulation AIDD t '+ str(i) +' number '+str(pnum)+'.csv',index_col=0)



def sens_test(df,simulation_num,leng):
    sdf = pd.DataFrame()
    sens = 0
    sign = 0
    for i in range(leng):
        if sens == 0 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 < df[str(simulation_num) + 'b'].iloc[i]*100:
            sens = 1
        if sens == 1 and 0 < df[str(simulation_num) + 't'].iloc[i]*100 and 0 > df[str(simulation_num) + 'b'].iloc[i]*100:
            sign = i
            break
    if sens == 1:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        if sign == 0:
            sdf.loc['sign','value'] = 194
        else:
            sdf.loc['sign','value'] = sign
        return sdf
    else:
        sdf.loc['vt','value'] = df[str(simulation_num) + 't'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vm','value'] = df[str(simulation_num) + 'm'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['vb','value'] = df[str(simulation_num) + 'b'][df[str(simulation_num) + 'm'].idxmax()]*100
        sdf.loc['d','value'] = df[str(simulation_num) + 'm'].idxmax()
        sdf.loc['sign','value'] = -1
        return sdf

sendf = {}


dnlist = [5,10,15,20,25,30,35,40,42,45,50]
for pnum in [2,3,4,5]:
    for i in dnlist:
        if i != 8:
            sendf[f'wdrsdf{i}_{pnum}'] = sens_test(dataframes['dr'+str(i)+'_'+str(pnum)], 2, len(dataframes['dr'+str(i)+'_'+str(pnum)]))
        else:
            sendf[f'wdrsdf{i}_{pnum}'] = sens_test(dataframes['dr'+str(i)+'_'+str(pnum)], 2, len(dataframes['dr'+str(i)+'_'+str(pnum)]))

namelist = [f'wdrsdf{i}_{pnum}' for i in dnlist for pnum in [2,3,4,5]]

wdrsdf = pd.concat([sendf[f'wdrsdf{i}_{pnum}'] for i in dnlist for pnum in [2,3,4,5]],axis=1)
wdrsdf.columns = namelist
wdrsdf = wdrsdf.T


wnlist = [5,8,10,15,20,25,30,35,40,45,50]

for pnum in [2,3,4,5]:
    for i in wnlist:
        if i != 8:
            sendf[f'wwfsdf{i}_{pnum}'] = sens_test(dataframes['wf'+str(i)+'_'+str(pnum)]/100, 2, len(dataframes['wf'+str(i)+'_'+str(pnum)]))
        else:
            sendf[f'wwfsdf{i}_{pnum}'] = sens_test(dataframes['wf'+str(i)+'_'+str(pnum)]/100, 2, len(dataframes['wf'+str(i)+'_'+str(pnum)]))

namelist = [f'wwfsdf{i}_{pnum}' for i in wnlist for pnum in [2,3,4,5]]

wwfsdf = pd.concat([sendf[f'wwfsdf{i}_{pnum}'] for i in wnlist for pnum in [2,3,4,5]],axis=1)
wwfsdf.columns = namelist

wwfsdf = wwfsdf.T



from matplotlib.legend_handler import HandlerTuple


factors = [2,3,4,5]
markers = ['o','v','^','s']
fig,ax = plt.subplots(figsize=(5,5),dpi=400)

for k in dnlist:
    if k != 42:
        tmp = wdrsdf[wdrsdf.index.str.contains('wdrsdf'+str(k)+'_')]
        a, b = np.polyfit(tmp['vm'], tmp['d'], 1)
        plt.plot(tmp['vm'], a * tmp['vm'] + b, color='orange', linestyle='dotted', alpha=0.7)
        for i in factors:
            tmp = wdrsdf[wdrsdf.index.str.contains('wdrsdf'+str(k)+'_'+str(i))]
            if float(tmp['sign'].values[0]) <0:
                plt.scatter(tmp['vm'], tmp['d'], color='orange', marker=markers[i - 2], alpha=0.7,facecolors='none')
            else:
                plt.scatter(tmp['vm'], tmp['d'], color='orange', marker=markers[i - 2], alpha=0.7)

    else:
        tmp = wdrsdf[wdrsdf.index.str.contains('wdrsdf'+str(k)+'_')]
        a,b = np.polyfit(tmp['vm'],tmp['d'],1)
        plt.plot(tmp['vm'],a*tmp['vm']+b,color='black',linestyle='-', alpha=1)
        for i in factors:
            tmp = wdrsdf[wdrsdf.index.str.contains('wdrsdf'+str(k)+'_'+str(i))].reset_index()
            if float(tmp['sign'].values[0]) <0:
                plt.scatter(tmp['vm'], tmp['d'], color='k', marker=markers[i - 2], alpha=1,facecolors='none')
            else:
                plt.scatter(tmp['vm'], tmp['d'], color='k', marker=markers[i - 2], alpha=1)
for k in wnlist:
    if k != 8:
        tmp = wwfsdf[wwfsdf.index.str.contains('wwfsdf'+str(k)+'_')]
        a, b = np.polyfit(tmp['vm'], tmp['d'], 1)
        plt.plot(tmp['vm'], a * tmp['vm'] + b, color='red', linestyle='dotted', alpha=0.7)
        for i in factors:
            tmp = wwfsdf[wwfsdf.index.str.contains('wwfsdf'+str(k)+'_'+str(i))]
            if float(tmp['sign'].values[0]) <0:
                plt.scatter(tmp['vm'], tmp['d'], color='red', marker=markers[i - 2], alpha=0.7,facecolors='none')
            else:
                plt.scatter(tmp['vm'], tmp['d'], color='red', marker=markers[i - 2], alpha=0.7)
    else:
        tmp = wwfsdf[wwfsdf.index.str.contains('wwfsdf'+str(k)+'_')]
        a,b = np.polyfit(tmp['vm'],tmp['d'],1)
        plt.plot(tmp['vm'],a*tmp['vm']+b,color='black',linestyle='-', alpha=1)
        for i in factors:
            tmp = wwfsdf[wwfsdf.index.str.contains('wwfsdf'+str(k)+'_'+str(i))]
            if float(tmp['sign'].values[0]) <0:
                plt.scatter(tmp['vm'], tmp['d'], color='k', marker=markers[i - 2], alpha=1,facecolors='none')
            else:
                plt.scatter(tmp['vm'], tmp['d'], color='k', marker=markers[i - 2], alpha=1)

for i in factors:
    plt.scatter(100, 100, color='k', marker=markers[i - 2],label = 'Factor '+str(i))
p1 =ax.scatter(100, 100, color='k', marker=markers[0])
p2 =ax.scatter(100, 100, color='k', marker=markers[1])
p3 =ax.scatter(100, 100, color='k', marker=markers[2])
p4 =ax.scatter(100, 100, color='k', marker=markers[3])
p5 =ax.scatter(100, 100, color='k', marker=markers[0],facecolors='none')
p6 =ax.scatter(100, 100, color='k', marker=markers[1],facecolors='none')
p7 =ax.scatter(100, 100, color='k', marker=markers[2],facecolors='none')
p8 =ax.scatter(100, 100, color='k', marker=markers[3],facecolors='none')

plt.axhline(y=10,linestyle='--',color='k',linewidth=0.5)
plt.axhline(y=15,linestyle='--',color='k',linewidth=0.5)
plt.axhline(y=20,linestyle='--',color='k',linewidth=0.5)
plt.axhline(y=25,linestyle='--',color='k',linewidth=0.5)
plt.axhline(y=30,linestyle='--',color='k',linewidth=0.5)
l = ax.legend([(p1, p5), (p2, p6), (p3, p7), (p4, p8)], ['Factor 2', 'Factor 3','Factor 4', 'Factor 5'],
               handler_map={tuple: HandlerTuple(ndivide=None)},loc='upper right'
              ,framealpha=1,edgecolor='black',prop = { "size": 10 })
plt.xlim([4,35])
plt.ylim([7,35])
plt.show()


